import pandas as pd
import numpy as np
import os
from sklearn.linear_model import Ridge
from sklearn.preprocessing import PolynomialFeatures
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


def main():
    if not os.path.exists("images"):
        os.makedirs("images")
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'KaiTi', 'SimSun']
    plt.rcParams['axes.unicode_minus'] = False

    # 读取2022-2024年数据
    df_2022_2024 = pd.read_csv("csv/dalian_weather_2022-2024_full.csv")
    df_2022_2024['日期'] = pd.to_datetime(df_2022_2024['日期'])

    # 读取2025年数据
    df_2025 = pd.read_csv("csv/dalian_weather_2025.csv")
    df_2025['日期'] = pd.to_datetime(df_2025['日期'])

    # 创建年份和月份特征
    df_2022_2024['年份'] = df_2022_2024['日期'].dt.year
    df_2022_2024['月份'] = df_2022_2024['日期'].dt.month

    # 计算2022-2024年每月平均最高气温
    monthly_data = df_2022_2024.groupby(['年份', '月份']).agg({
        '最高气温': 'mean'
    }).reset_index()

    # 添加时间索引
    monthly_data['时间索引'] = np.arange(len(monthly_data))
    monthly_data['月份日期'] = pd.to_datetime(
        monthly_data['年份'].astype(str) + '-' + monthly_data['月份'].astype(str) + '-01')

    # 周期性编码月份
    monthly_data['月份_sin'] = np.sin(2 * np.pi * monthly_data['月份'] / 12)
    monthly_data['月份_cos'] = np.cos(2 * np.pi * monthly_data['月份'] / 12)

    # 训练集：2022-2024年的所有月份数据
    train = monthly_data.copy()
    # 特征工程
    features = ['时间索引', '月份_sin', '月份_cos']
    X_train = train[features]
    y_train = train['最高气温']

    # 创建2025年 1-6月的预测数据
    future_months = pd.DataFrame({
        '年份': [2025] * 6,
        '月份': range(1, 7),
        '时间索引': np.arange(len(monthly_data), len(monthly_data) + 6)
    })

    # 添加周期性特征
    future_months['月份_sin'] = np.sin(2 * np.pi * future_months['月份'] / 12)
    future_months['月份_cos'] = np.cos(2 * np.pi * future_months['月份'] / 12)

    future_months['月份日期'] = pd.to_datetime(
        future_months['年份'].astype(str) + '-' + future_months['月份'].astype(str) + '-01')

    # 计算2025年1-6月的实际月平均气温
    df_2025 = df_2025[df_2025['日期'] <= '2025-06-30']
    df_2025['月份'] = df_2025['日期'].dt.month
    monthly_2025 = df_2025.groupby(['月份']).agg({
        '最高气温': 'mean'
    }).reset_index()
    monthly_2025.rename(columns={'最高气温': '实际最高气温'}, inplace=True)
    monthly_2025['年份'] = 2025
    monthly_2025['月份日期'] = pd.to_datetime('2025-' + monthly_2025['月份'].astype(str) + '-01')

    # 合并预测数据和实际数据
    future_months = pd.merge(future_months, monthly_2025[['月份日期', '实际最高气温']],
                             on='月份日期', how='left')

    # 准备不同多项式次数的模型
    degrees = [2, 3, 4]
    models = {}
    predictions = {}

    # 训练不同次数的多项式模型
    for degree in degrees:
        # 创建多项式特征
        poly = PolynomialFeatures(degree=degree, include_bias=False)
        X_train_poly = poly.fit_transform(X_train)

        # 训练模型
        model = Ridge(alpha=0.1)
        model.fit(X_train_poly, y_train)
        models[degree] = (poly, model)

        # 进行预测
        X_future = future_months[features]
        X_future_poly = poly.transform(X_future)
        future_predictions = model.predict(X_future_poly)
        predictions[degree] = future_predictions
        future_months[f'预测最高气温_{degree}次'] = future_predictions

    # 打印预测结果和实际数据
    print("2025年1-6月不同多项式拟合预测与实际最高气温(月平均):")
    print(future_months[['年份', '月份', '实际最高气温',
                         '预测最高气温_2次', '预测最高气温_3次', '预测最高气温_4次']].to_string(index=False))

    # 可视化所有数据
    plt.figure(figsize=(16, 10))

    # 绘制2022-2024年月平均数据
    plt.plot(monthly_data['月份日期'], monthly_data['最高气温'], 'bo-',
             markersize=8, linewidth=2, label='2022-2024月平均')

    # 绘制2025年实际月平均数据
    plt.plot(monthly_2025['月份日期'], monthly_2025['实际最高气温'], 'mo-',
             markersize=8, linewidth=2, label='2025实际月平均')

    # 绘制不同多项式次数的预测数据
    colors = ['green', 'orange', 'red']
    line_styles = ['--', '-.', ':']
    for i, degree in enumerate(degrees):
        plt.plot(future_months['月份日期'], predictions[degree],
                 color=colors[i], linestyle=line_styles[i],
                 markersize=10, linewidth=2,
                 label=f'{degree}次多项式预测')

    # 添加标签和标题
    plt.title('大连最高气温分析 - 不同多项式拟合比较 (2022-2024历史数据与2025预测/实际对比)', fontsize=16)
    plt.xlabel('日期', fontsize=14)
    plt.ylabel('最高气温(℃)', fontsize=14)
    plt.legend(loc='best', fontsize=12)

    # 设置x轴格式 - 每月一个刻度
    plt.gca().xaxis.set_major_locator(mdates.MonthLocator(interval=1))
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    plt.xticks(rotation=45, fontsize=12)
    plt.yticks(fontsize=12)

    # 设置网格
    plt.grid(True, linestyle='--', alpha=0.7)

    # 设置日期范围
    min_date = min(monthly_data['月份日期'].min(), future_months['月份日期'].min())
    max_date = max(future_months['月份日期'].max(), monthly_2025['月份日期'].max())
    plt.xlim(min_date, max_date)

    # 添加图例说明
    plt.figtext(0.5, 0.01, "注: 展示了2次、3次和4次多项式拟合的预测结果对比",
                ha="center", fontsize=12, bbox={"facecolor": "orange", "alpha": 0.2, "pad": 5})

    plt.tight_layout()
    plt.subplots_adjust(bottom=0.15)  # 为底部文本留出空间
    plt.show()


if __name__ == '__main__':
    main()